%%time
import sys, os

# get current directory
path = os.getcwd()

# get parent directory
parent_directory = os.path.sep.join(path.split(os.path.sep)[:-4])

# add utils folder to current working path
sys.path.append(parent_directory+"/subfunctions/utils")

# add integration folder to current working path
sys.path.append(parent_directory+"/subfunctions/integration")

# add TRA folder to current working path
sys.path.append(parent_directory+"/demos/AdvectiveBarriers/TRA")

%%time
import scipy.io as sio

#Import velocity data from file in data-folder
mat_file = sio.loadmat('../../../../data/Aviso/AVISO.mat')

U = mat_file['u']
V = mat_file['v']
x = mat_file['x']
y = mat_file['y']
time_data = mat_file['t']

%%time
# Initial time (in days)
t0 = 0

# Final time (in days)
tN = 25

# time step-size (in days)
dt = .1

time = np.arange(t0, tN+dt, dt)

# length of time interval (in days)
lenT = tN-t0

# longitudinal and latitudinal boundaries (in degrees)
xmin = -3
xmax = 1
ymin = -32
ymax = -24

# make sure that domain is in the data
assert (xmax <= np.max(X) and xmin >= np.min(X) and ymin >= np.min(Y) and ymax <= np.max(Y) and t0 >= np.min(time_data) and tN <= np.max(time_data)),"The domains you are chooising are outside the domain of the data!!!!! --> redefine spatial/temporal domain"
    
# spacing of meshgrid (in degrees)
dx = 0.02
dy = 0.02

x_domain = np.arange(xmin, xmax + dx, dx)
y_domain = np.arange(ymin, ymax + dy, dy)

X_domain, Y_domain = np.meshgrid(x_domain, y_domain)

# Size of arrays
Nx = X_domain.shape[1]
Ny = Y_domain.shape[0]

%%time
# Import interpolation function for unsteady flow field
from ipynb.fs.defs.Interpolant import interpolant_unsteady

# set nan values to zero so that we can apply interpolant. Interpolant does not work if the array contains nan values
U[np.isnan(U)] = 0
V[np.isnan(V)] = 0

# Interpolate velocity data using cubic spatial interpolation
Interpolant = interpolant_unsteady(X, Y, U, V, time, method = "cubic")

%%time
# Import package for computing trajectories/velocity along trajectories
from ipynb.fs.defs.integration_dFdt import integration_dFdt

# Import package which checks particle location
from ipynb.fs.defs.check_location import check_location

# Import package for progress bar
from tqdm.notebook import tqdm

# Import package for parallel computing
from joblib import Parallel, delayed

# Import package for TRA
from ipynb.fs.defs.TRA import _TRA

# set v0 is equal to the average velocity in the Agulhas region
v0 = np.mean(np.sqrt(U.ravel()**2+V.ravel()**2))

def parallel_TRA(i):
    
    TRA_parallel = X_domain[0,:].copy()*np.nan
    
    for j in range(X_domain.shape[1]):
        
        # set initial condition
        x = np.array([X_domain[i, j], Y_domain[i, j]])
        
        # only compute TRA for trajectories starting in region where velocity field is defined
        if check_location(X, Y, defined_domain, x)[0] == "IN":

            # compute velocity along trajectory
            dFdt = integration_dFdt(time, x, X, Y, Interpolant, periodic, defined_domain, bool_unsteady, dt_data, delta)[1]
    
            TRA_parallel[j] = _TRA(lenT, dFdt, v0)
        
        # else set value to nan
        else:
            
            TRA_parallel[j] = np.nan
        
    return TRA_parallel

TRA = np.array(Parallel(n_jobs=Ncores, verbose = 0)(delayed(parallel_TRA)(i) for i in tqdm(range(X_domain.shape[0]))))

######################## PLOT RESULTS ########################
import matplotlib.pyplot as plt
import matplotlib.colors as colors

# Figure/Axis
fig = plt.figure(figsize=(9, 6), dpi = 600)
ax = plt.axes()

# Contourplot of TRA over meshgrid of initial conditions
cax = ax.contourf(X_domain, Y_domain, np.ma.masked_invalid(TRA), levels = 600, cmap = "gist_rainbow_r")

# Axis Labels
ax.set_xlabel("long (°)", fontsize = 16)
ax.set_ylabel("lat (°)", fontsize = 16)

# Ticks
ax.set_xticks(np.arange(np.min(X_domain), np.max(X_domain), 1))
ax.set_yticks(np.arange(np.min(Y_domain), np.max(Y_domain), 1))

# Colorbar
cbar = fig.colorbar(cax, ticks = np.linspace(0, 0.5, 11))
cbar.ax.set_ylabel(r'$(\dfrac{rad}{d})$', rotation = 0, labelpad = 20, fontsize = 10)

# Title
ax.set_title(r'$ \mathrm{\overline{TRA}}$'+f'$_{{{int(time[0])}}}^{{{int(time[-1])}}}$', fontsize = 20)

plt.show()

# import random number generator
import random

# 10% subsampling
rdn10 = random.sample(range(0, Nx*Ny), int(Nx*Ny/10))

# subsample initial grid and TRA field (=subsample trajectories)
x_subsampled10 = X_domain.ravel()[rdn10]
y_subsampled10 = Y_domain.ravel()[rdn10]
TRA_subsampled10 = TRA.ravel()[rdn10]

# 1% subsampling
rdn100 = random.sample(range(0, Nx*Ny), int(Nx*Ny/100))

# subsample initial grid and TRA field (=subsample trajectories)
x_subsampled100 = X_domain.ravel()[rdn100]
y_subsampled100 = Y_domain.ravel()[rdn100]
TRA_subsampled100 = TRA.ravel()[rdn100]

# 0.1% subsampling
rdn1000 = random.sample(range(0, Nx*Ny), int(Nx*Ny/1000))

# subsample initial grid and TRA field (=subsample trajectories)
x_subsampled1000 = X_domain.ravel()[rdn1000]
y_subsampled1000 = Y_domain.ravel()[rdn1000]
TRA_subsampled1000 = TRA.ravel()[rdn1000]

# Import interpolant for scatter interpolation
from scipy.interpolate import Rbf

# compute interpolation object
interp10 = Rbf(x_subsampled10, y_subsampled10, TRA_subsampled10, function = "linear")
interp100 = Rbf(x_subsampled100, y_subsampled100, TRA_subsampled100, function = "linear")
interp1000 = Rbf(x_subsampled1000, y_subsampled1000, TRA_subsampled1000, function = "linear")

# re-evaluate interpolant over original (initial) meshgrid
TRA_interpolated10 = interp10(X_domain, Y_domain) 
TRA_interpolated100 = interp100(X_domain, Y_domain) 
TRA_interpolated1000 = interp1000(X_domain, Y_domain) 

# Figure/Axis
fig, (ax0, ax1, ax2) = plt.subplots(nrows=1, ncols=3, dpi = 600, figsize=(12, 4))

# Contourplot of TRA over meshgrid of initial conditions
cax = ax0.contourf(X_domain, Y_domain, np.ma.masked_invalid(TRA_interpolated10), cmap = "gist_rainbow_r", levels = 600)
cax = ax1.contourf(X_domain, Y_domain, np.ma.masked_invalid(TRA_interpolated100), cmap = "gist_rainbow_r", levels = 600)
cax = ax2.contourf(X_domain, Y_domain, np.ma.masked_invalid(TRA_interpolated1000), cmap = "gist_rainbow_r", levels = 600)

# Axis Labels
ax0.set_xlabel("long (°)", fontsize = 16)
ax1.set_xlabel("long (°)", fontsize = 16)
ax2.set_xlabel("long (°)", fontsize = 16)
ax0.set_ylabel("lat (°)", fontsize = 16)

ax0.set_title("10% subsampling")
ax1.set_title("1% subsampling")
ax2.set_title("0.1% subsampling")

# Colorbar
cbar = fig.colorbar(cax, ticks = np.linspace(0, 0.4, 11))
cbar.ax.set_ylabel(r'$(\dfrac{rad}{d})$', rotation = 0, labelpad = 20, fontsize = 10)

# Title
ax.set_title(r'$ \mathrm{\overline{TRA}}$'+f'$_{{{int(time[0])}}}^{{{int(time[-1])}}}$', fontsize = 20)

plt.show()